# -*- coding: utf-8 -*-
#Created on Mon Aug 6 9:49:15 2018
#@author: JockJo

from pymongo import MongoClient


class Mongo():
    mongo_collection = 'query_{}'
    mongo_db = 'news_query'
    client=MongoClient()

    def __init__(self, client=None, port=None, mongo_collection=None, mongo_db=None):
        if mongo_collection is not None:
            self.mongo_collection = mongo_collection
        if mongo_db is not None:
            self.mongo_db = mongo_db

        if client is None:
            self.client = MongoClient('localhost', 27017)     # default connection
        else:
            self.client = MongoClient(client, port)            # specific connection
        self.db = self.client[self.mongo_db]

    def get_item(self, collection_name, item, field):
        record = self.db[self.mongo_collection.format(collection_name)].find_one({'field': item})
        if record:
            return not None
        else:
            return None

    def get_company(self, collection_name):
        company = self.db[self.mongo_collection.format(collection_name)].find()
        company_list = []
        for c in company:
            if c['Conme'] is not '' and c['Stknme'] is not '':
                if type(c['Stkcd']) is type(''):
                    query_symbol = 'sh'+c['Stkcd'] if c['Stkcd'][0] == '6' and c['Stkcd'][1] == '0' else 'sz'+c['Stkcd']
                else:
                    query_symbol = 'sz'+str(c['Stkcd']) if c['Stkcd'] < 600000 else 'sh'+str(c['Stkcd'])
                company_list.append([c['Stknme'], c['Conme'], c['Stkcd'],query_symbol])  # 证券名称,公司全称,证券代码
        return company_list

    def process_item(self, item, collection_name):
        record = {'Stkcd':item['Stkcd'],'Stknme':item['Stknme'],
                  'Conme':item['Conme'], 'Conme_en':item['Conme_en'],
                  'Estbdt':item['Estbdt'], 'Listdt':item['Listdt'],
                  'Markettype':item['Markettype']}
        self.db[self.mongo_collection.format(collection_name)].update({'_id':item['Stkcd']}, {'$set':record}, upsert=True)
        return item

    def close(self):
        self.client.close()